using System;
using System.Collections.Generic;
using System.Text;
using System.IO;
using System.Runtime.InteropServices;

namespace Brett
{
    public class VirusUtilities
    {
        /// <summary>
        /// Used to get the pid of a process that has a handle on a file.
        /// This method uses SysInternals Handle utility. It will throw an exception if handle.exe
        /// is not in the same directory as this library.
        /// </summary>
        /// <param name="file"></param>
        /// <returns></returns>
        public static int GetProcessWithHandle(string file)
        {
            if (!System.IO.File.Exists("handle.exe"))
            {
                throw new FileNotFoundException("Could not find handle.exe. Make sure it is in the current directory. It can be downloaded from http://technet.microsoft.com/en-us/sysinternals/bb896655.aspx", "handle.exe");
            }
            using (System.Diagnostics.Process handleProcess = new System.Diagnostics.Process())
            {
                handleProcess.StartInfo.FileName = "handle.exe";
                handleProcess.StartInfo.Arguments = '"' + file + '"';
                handleProcess.StartInfo.RedirectStandardOutput = true;
                handleProcess.StartInfo.UseShellExecute = false;
                handleProcess.Start();
                string output = handleProcess.StandardOutput.ReadToEnd();
                if (!output.Contains("pid: "))
                {
                    return -1;
                }
                int i = output.IndexOf("pid: ") + 5;
                string pid = string.Empty;
                while (true)
                {
                    char c = output[i];
                    if (c.Equals(' '))
                    {
                        break;
                    }
                    pid += c;
                    i++;
                }
                return int.Parse(pid);
            }
        }
        
        public static byte[] ReadFromLockedFile(string file)
        {
            int pid = GetProcessWithHandle(file);
            if (!pid.Equals(-1))
            {
                try
                {
                    System.Diagnostics.Process process = System.Diagnostics.Process.GetProcessById(pid);
                    process.Kill();
                    process.WaitForExit();
                    return File.ReadAllBytes(file);
                }
                catch
                {
                    return null;
                }
            }
            return null;
        }

        public static Dictionary<string, byte[]> ReadMemoryFromProcess(int processID)
        {
            Dictionary<string, byte[]> moduleMemory = new Dictionary<string, byte[]>();
            
            System.Diagnostics.Process process = System.Diagnostics.Process.GetProcessById(processID);
            byte[] temp;
            foreach (System.Diagnostics.ProcessModule pm in process.Modules)
            { 
                temp = new byte[pm.ModuleMemorySize];
                int read;
                if (ReadProcessMemory(process.Handle, pm.BaseAddress, temp, temp.Length, out read))
                {
                    moduleMemory.Add(pm.ModuleName, temp);
                }
                temp = null;
            }
            
            return moduleMemory;
        }

        [DllImport("kernel32.dll", SetLastError = true)]
        static extern bool ReadProcessMemory(
          IntPtr hProcess,
          IntPtr lpBaseAddress,
          [Out()] byte[] lpBuffer,
          int dwSize,
          out int lpNumberOfBytesRead
         );
        
        private void Test()
        {
            byte[] data = ReadFromLockedFile(@"c:\windows\piserver.exe");
        }
    }
}
